from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex


class MyAES:
    """
        注意，AES 源码里限定如下：
        # Size of a data block (in bytes)
        block_size = 16
        # Size of a key (in bytes)
        key_size = (16, 24, 32)
        所以对skey的size要进行规范化
    """
    def __init__(self,skey:str):
        # print(skey)
        # 取 key_size = 16
        assert len(skey)>=16
        self.skey = skey[:16] #前16个字符,要求字符为ascii字符
        self._aes = AEScrypto()

    def encrypt(self,text:str)->bytes:
        # text = self._aes.add_to_16(text)
        ans = self._aes.encrypt(text,self.skey)
        return ans

    def decrypt(self,cipher_bytes:bytes)->str:
        ans = self._aes.decrypt(cipher_bytes,self.skey)
        ans = ans.strip('\0')
        return ans

class AEScrypto:

    def __init__(self):
        self.iv = b'qqqqqqqqqqqqqqqq'#这个是位移

    # 如果text不足16位的倍数就用空格补足为16位
    def add_to_16(self,text):
        if len(text.encode('utf-8')) % 16:
            add = 16 - (len(text.encode('utf-8')) % 16)
        else:
            add = 0
        text = text + ('\0' * add)
        return text.encode('utf-8')

    # 加密函数
    def encrypt(self,text,skey):
        key = skey.encode('utf-8')
        mode = AES.MODE_CBC
        text = self.add_to_16(text)
        cryptos = AES.new(key, mode, self.iv)
        cipher_text = cryptos.encrypt(text)
        # 因为AES加密后的字符串不一定是ascii字符集的，输出保存可能存在问题，所以这里转为16进制字符串
        return b2a_hex(cipher_text)

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt(self,text, skey):
        key = skey.encode('utf-8')
        mode = AES.MODE_CBC
        cryptos = AES.new(key, mode, self.iv)
        plain_text = cryptos.decrypt(a2b_hex(text))
        # return bytes.decode(plain_text).rstrip('\0')
        ans = plain_text.rstrip(b'\0').decode("utf-8")
        return ans

if __name__ == '__main__':
    
    skey = '9999999999999999'			#秘钥
    text = 'Rukia'				#被加密内容
    ec = AEScrypto()
    encrypted_text = ec.encrypt(text, skey)  # 加密
    decrypted_text = ec.decrypt(encrypted_text, skey)  # 解密
    print("加密:", encrypted_text)
    print("解密:", decrypted_text)